from typing import overload, Union

@overload
def baz(a: int, b) -> int: ...
@overload
def baz(a: str, b) -> str: ...
def baz(a: Union[int, str], b) -> Union[int, str]: ...
